from paramiko.client import MissingHostKeyPolicy, SSHClient, AutoAddPolicy
from may import Server


class SSHConnection(SSHClient):
    def __init__(self, pool, server: Server, missing_host_key_policy: MissingHostKeyPolicy) -> None:
        super().__init__()
        self.pool = pool
        self.server = server
        self.set_missing_host_key_policy(missing_host_key_policy)
        self.connect(
            server.host,
            server.port if server.port else pool.config["default_port"],
            **server.user.get_auth_params(),
            timeout=self.pool.config["connect_timeout"]
        )

    def __exit__(self, type, value, traceback):
        """
        support 'with' grammar
        ssh_conn_pool: SSHConnectionPool

        with ssh_conn_pool.get_conn(hostname) as conn:
            stdin, stdout, stderr = conn.exec_command('ls -l')
        """
        self.pool.add(self)
    
    close = __exit__

    def destroy(self):
        "instead of close func to do close"
        super().close()


class SSHConnectionPool():
    """ssh connection pool"""

    def __init__(self, servers: list, config: dict) -> None:
        self.min_num_per_server = config["min_num_per_server"]
        self.max_num_per_server = config["max_num_per_server"]
        self._pool = {}
        self._conn_num = {}
        self._host2server = {}
        self.__missing_host_key_policy = AutoAddPolicy()
        for server in servers:
            self._pool[server.host] = [
                self._create_conn(server)
                for _ in range(self.min_num_per_server)
            ]
            self._conn_num[server.host] = len(self._pool[server.host])
            self._host2server[server.host] = server

    def _create_conn(self, server):
        return SSHConnection(self, server, self.__missing_host_key_policy)

    def add(self, conn: SSHConnection):
        self._pool[conn.server.host].append(conn)

    def pop(self, host):
        assert host in self._pool, "host: {} had not configured".format(host)
        if not self._pool[host] and self._conn_num[host] < self.max_num_per_server:
            self._conn_num[host] += 1
            new_conn = self._create_conn(self._host2server[host])
            self._pool[host].append(new_conn)
        return self._pool[host].pop()

    get_conn = pop
